import tkinter as tk
from tkinter import filedialog, messagebox
from cryptography.fernet import Fernet

# 生成密钥
def generate_key():
    key = Fernet.generate_key()
    with open("secret.key", "wb") as key_file:
        key_file.write(key)
    messagebox.showinfo("成功", "密钥已生成并保存到 secret.key 文件中")

# 加载密钥
def load_key():
    return open("secret.key", "rb").read()

# 加密文件
def encrypt_file(file_path, key):
    fernet = Fernet(key)
    with open(file_path, "rb") as file:
        original = file.read()
    encrypted = fernet.encrypt(original)
    with open(file_path, "wb") as encrypted_file:
        encrypted_file.write(encrypted)
    messagebox.showinfo("成功", f"文件 {file_path} 已加密")

# 解密文件
def decrypt_file(file_path, key):
    fernet = Fernet(key)
    with open(file_path, "rb") as encrypted_file:
        encrypted = encrypted_file.read()
    decrypted = fernet.decrypt(encrypted)
    with open(file_path, "wb") as decrypted_file:
        decrypted_file.write(decrypted)
    messagebox.showinfo("成功", f"文件 {file_path} 已解密")

# 选择文件
def select_file():
    file_path = filedialog.askopenfilename()
    if file_path:
        return file_path
    return None

# 创建主窗口
def create_main_window():
    root = tk.Tk()
    root.title("文件加密/解密工具")

    # 生成密钥按钮
    generate_key_button = tk.Button(root, text="生成密钥", command=generate_key)
    generate_key_button.pack(pady=10)

    # 选择文件按钮
    select_file_button = tk.Button(root, text="选择文件", command=lambda: select_file())
    select_file_button.pack(pady=10)

    # 加密文件按钮
    encrypt_button = tk.Button(root, text="加密文件", command=lambda: encrypt_file(select_file(), load_key()))
    encrypt_button.pack(pady=10)

    # 解密文件按钮
    decrypt_button = tk.Button(root, text="解密文件", command=lambda: decrypt_file(select_file(), load_key()))
    decrypt_button.pack(pady=10)

    root.mainloop()

if __name__ == "__main__":
    create_main_window()